{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "3779ad61",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import io\n",
    "import re\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score\n",
    "from sklearn.metrics import classification_report\n",
    "from sklearn.preprocessing import MinMaxScaler, StandardScaler,LabelEncoder\n",
    "from sklearn.decomposition import PCA\n",
    "import sqlite3\n",
    "from prophet import Prophet\n",
    "from prophet.plot import add_changepoints_to_plot\n",
    "import datetime as dt\n",
    "from sklearn import linear_model\n",
    "import xgboost as xgb\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from imblearn.over_sampling import SMOTE\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import logging\n",
    "import warnings\n",
    "import json\n",
    "%matplotlib inline\n",
    "warnings.filterwarnings('ignore')\n",
    "pd.set_option('display.max_rows', None)\n",
    "pd.set_option('display.max_columns', 500)\n",
    "pd.set_option('display.width', 1000)\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.CRITICAL)\n",
    "import openai\n",
    "import pandas as pd\n",
    "import sqlite3\n",
    "from openai.api_resources.completion import Completion\n",
    "import sqlparse\n",
    "import time\n",
    "from sqlparse.sql import IdentifierList, Identifier\n",
    "from sqlparse.tokens import Token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "28442f03",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Index(['OBJECTID', 'FOD_ID', 'FPA_ID', 'SOURCE_SYSTEM_TYPE', 'SOURCE_SYSTEM', 'NWCG_REPORTING_AGENCY', 'NWCG_REPORTING_UNIT_ID', 'NWCG_REPORTING_UNIT_NAME', 'SOURCE_REPORTING_UNIT', 'SOURCE_REPORTING_UNIT_NAME', 'LOCAL_FIRE_REPORT_ID', 'LOCAL_INCIDENT_ID', 'FIRE_CODE', 'FIRE_NAME', 'ICS_209_INCIDENT_NUMBER', 'ICS_209_NAME', 'MTBS_ID', 'MTBS_FIRE_NAME', 'COMPLEX_NAME', 'FIRE_YEAR', 'DISCOVERY_DATE', 'DISCOVERY_DOY', 'DISCOVERY_TIME', 'STAT_CAUSE_CODE', 'STAT_CAUSE_DESCR', 'CONT_DATE', 'CONT_DOY', 'CONT_TIME', 'FIRE_SIZE', 'FIRE_SIZE_CLASS', 'LATITUDE', 'LONGITUDE', 'OWNER_CODE', 'OWNER_DESCR', 'STATE', 'COUNTY', 'FIPS_CODE', 'FIPS_NAME', 'Shape'], dtype='object')\""
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cnx = sqlite3.connect('FPA_FOD_20170508.sqlite')\n",
    "sql = \"select * from fires\"\n",
    "\n",
    "df = pd.read_sql_query(sql, cnx)\n",
    "df.head()\n",
    "str(df.columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2b37d637",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>OBJECTID</th>\n",
       "      <th>FOD_ID</th>\n",
       "      <th>FPA_ID</th>\n",
       "      <th>SOURCE_SYSTEM_TYPE</th>\n",
       "      <th>SOURCE_SYSTEM</th>\n",
       "      <th>NWCG_REPORTING_AGENCY</th>\n",
       "      <th>NWCG_REPORTING_UNIT_ID</th>\n",
       "      <th>NWCG_REPORTING_UNIT_NAME</th>\n",
       "      <th>SOURCE_REPORTING_UNIT</th>\n",
       "      <th>SOURCE_REPORTING_UNIT_NAME</th>\n",
       "      <th>LOCAL_FIRE_REPORT_ID</th>\n",
       "      <th>LOCAL_INCIDENT_ID</th>\n",
       "      <th>FIRE_CODE</th>\n",
       "      <th>FIRE_NAME</th>\n",
       "      <th>ICS_209_INCIDENT_NUMBER</th>\n",
       "      <th>ICS_209_NAME</th>\n",
       "      <th>MTBS_ID</th>\n",
       "      <th>MTBS_FIRE_NAME</th>\n",
       "      <th>COMPLEX_NAME</th>\n",
       "      <th>FIRE_YEAR</th>\n",
       "      <th>DISCOVERY_DATE</th>\n",
       "      <th>DISCOVERY_DOY</th>\n",
       "      <th>DISCOVERY_TIME</th>\n",
       "      <th>STAT_CAUSE_CODE</th>\n",
       "      <th>STAT_CAUSE_DESCR</th>\n",
       "      <th>CONT_DATE</th>\n",
       "      <th>CONT_DOY</th>\n",
       "      <th>CONT_TIME</th>\n",
       "      <th>FIRE_SIZE</th>\n",
       "      <th>FIRE_SIZE_CLASS</th>\n",
       "      <th>LATITUDE</th>\n",
       "      <th>LONGITUDE</th>\n",
       "      <th>OWNER_CODE</th>\n",
       "      <th>OWNER_DESCR</th>\n",
       "      <th>STATE</th>\n",
       "      <th>COUNTY</th>\n",
       "      <th>FIPS_CODE</th>\n",
       "      <th>FIPS_NAME</th>\n",
       "      <th>Shape</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>FS-1418826</td>\n",
       "      <td>FED</td>\n",
       "      <td>FS-FIRESTAT</td>\n",
       "      <td>FS</td>\n",
       "      <td>USCAPNF</td>\n",
       "      <td>Plumas National Forest</td>\n",
       "      <td>0511</td>\n",
       "      <td>Plumas National Forest</td>\n",
       "      <td>1</td>\n",
       "      <td>PNF-47</td>\n",
       "      <td>BJ8K</td>\n",
       "      <td>FOUNTAIN</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>2005</td>\n",
       "      <td>2453403.5</td>\n",
       "      <td>33</td>\n",
       "      <td>1300</td>\n",
       "      <td>9.0</td>\n",
       "      <td>Miscellaneous</td>\n",
       "      <td>2453403.5</td>\n",
       "      <td>33.0</td>\n",
       "      <td>1730</td>\n",
       "      <td>0.10</td>\n",
       "      <td>A</td>\n",
       "      <td>40.036944</td>\n",
       "      <td>-121.005833</td>\n",
       "      <td>5.0</td>\n",
       "      <td>USFS</td>\n",
       "      <td>CA</td>\n",
       "      <td>63</td>\n",
       "      <td>063</td>\n",
       "      <td>Plumas</td>\n",
       "      <td>b'\\x00\\x01\\xad\\x10\\x00\\x00\\xe8d\\xc2\\x92_@^\\xc0...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "      <td>FS-1418827</td>\n",
       "      <td>FED</td>\n",
       "      <td>FS-FIRESTAT</td>\n",
       "      <td>FS</td>\n",
       "      <td>USCAENF</td>\n",
       "      <td>Eldorado National Forest</td>\n",
       "      <td>0503</td>\n",
       "      <td>Eldorado National Forest</td>\n",
       "      <td>13</td>\n",
       "      <td>13</td>\n",
       "      <td>AAC0</td>\n",
       "      <td>PIGEON</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>2004</td>\n",
       "      <td>2453137.5</td>\n",
       "      <td>133</td>\n",
       "      <td>0845</td>\n",
       "      <td>1.0</td>\n",
       "      <td>Lightning</td>\n",
       "      <td>2453137.5</td>\n",
       "      <td>133.0</td>\n",
       "      <td>1530</td>\n",
       "      <td>0.25</td>\n",
       "      <td>A</td>\n",
       "      <td>38.933056</td>\n",
       "      <td>-120.404444</td>\n",
       "      <td>5.0</td>\n",
       "      <td>USFS</td>\n",
       "      <td>CA</td>\n",
       "      <td>61</td>\n",
       "      <td>061</td>\n",
       "      <td>Placer</td>\n",
       "      <td>b'\\x00\\x01\\xad\\x10\\x00\\x00T\\xb6\\xeej\\xe2\\x19^\\...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "      <td>FS-1418835</td>\n",
       "      <td>FED</td>\n",
       "      <td>FS-FIRESTAT</td>\n",
       "      <td>FS</td>\n",
       "      <td>USCAENF</td>\n",
       "      <td>Eldorado National Forest</td>\n",
       "      <td>0503</td>\n",
       "      <td>Eldorado National Forest</td>\n",
       "      <td>27</td>\n",
       "      <td>021</td>\n",
       "      <td>A32W</td>\n",
       "      <td>SLACK</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>2004</td>\n",
       "      <td>2453156.5</td>\n",
       "      <td>152</td>\n",
       "      <td>1921</td>\n",
       "      <td>5.0</td>\n",
       "      <td>Debris Burning</td>\n",
       "      <td>2453156.5</td>\n",
       "      <td>152.0</td>\n",
       "      <td>2024</td>\n",
       "      <td>0.10</td>\n",
       "      <td>A</td>\n",
       "      <td>38.984167</td>\n",
       "      <td>-120.735556</td>\n",
       "      <td>13.0</td>\n",
       "      <td>STATE OR PRIVATE</td>\n",
       "      <td>CA</td>\n",
       "      <td>17</td>\n",
       "      <td>017</td>\n",
       "      <td>El Dorado</td>\n",
       "      <td>b'\\x00\\x01\\xad\\x10\\x00\\x00\\xd0\\xa5\\xa0W\\x13/^\\...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>4</td>\n",
       "      <td>FS-1418845</td>\n",
       "      <td>FED</td>\n",
       "      <td>FS-FIRESTAT</td>\n",
       "      <td>FS</td>\n",
       "      <td>USCAENF</td>\n",
       "      <td>Eldorado National Forest</td>\n",
       "      <td>0503</td>\n",
       "      <td>Eldorado National Forest</td>\n",
       "      <td>43</td>\n",
       "      <td>6</td>\n",
       "      <td>None</td>\n",
       "      <td>DEER</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>2004</td>\n",
       "      <td>2453184.5</td>\n",
       "      <td>180</td>\n",
       "      <td>1600</td>\n",
       "      <td>1.0</td>\n",
       "      <td>Lightning</td>\n",
       "      <td>2453189.5</td>\n",
       "      <td>185.0</td>\n",
       "      <td>1400</td>\n",
       "      <td>0.10</td>\n",
       "      <td>A</td>\n",
       "      <td>38.559167</td>\n",
       "      <td>-119.913333</td>\n",
       "      <td>5.0</td>\n",
       "      <td>USFS</td>\n",
       "      <td>CA</td>\n",
       "      <td>3</td>\n",
       "      <td>003</td>\n",
       "      <td>Alpine</td>\n",
       "      <td>b'\\x00\\x01\\xad\\x10\\x00\\x00\\x94\\xac\\xa3\\rt\\xfa]...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>5</td>\n",
       "      <td>FS-1418847</td>\n",
       "      <td>FED</td>\n",
       "      <td>FS-FIRESTAT</td>\n",
       "      <td>FS</td>\n",
       "      <td>USCAENF</td>\n",
       "      <td>Eldorado National Forest</td>\n",
       "      <td>0503</td>\n",
       "      <td>Eldorado National Forest</td>\n",
       "      <td>44</td>\n",
       "      <td>7</td>\n",
       "      <td>None</td>\n",
       "      <td>STEVENOT</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>2004</td>\n",
       "      <td>2453184.5</td>\n",
       "      <td>180</td>\n",
       "      <td>1600</td>\n",
       "      <td>1.0</td>\n",
       "      <td>Lightning</td>\n",
       "      <td>2453189.5</td>\n",
       "      <td>185.0</td>\n",
       "      <td>1200</td>\n",
       "      <td>0.10</td>\n",
       "      <td>A</td>\n",
       "      <td>38.559167</td>\n",
       "      <td>-119.933056</td>\n",
       "      <td>5.0</td>\n",
       "      <td>USFS</td>\n",
       "      <td>CA</td>\n",
       "      <td>3</td>\n",
       "      <td>003</td>\n",
       "      <td>Alpine</td>\n",
       "      <td>b'\\x00\\x01\\xad\\x10\\x00\\x00@\\xe3\\xaa.\\xb7\\xfb]\\...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   OBJECTID  FOD_ID      FPA_ID SOURCE_SYSTEM_TYPE SOURCE_SYSTEM NWCG_REPORTING_AGENCY NWCG_REPORTING_UNIT_ID  NWCG_REPORTING_UNIT_NAME SOURCE_REPORTING_UNIT SOURCE_REPORTING_UNIT_NAME LOCAL_FIRE_REPORT_ID LOCAL_INCIDENT_ID FIRE_CODE FIRE_NAME ICS_209_INCIDENT_NUMBER ICS_209_NAME MTBS_ID MTBS_FIRE_NAME COMPLEX_NAME  FIRE_YEAR  DISCOVERY_DATE  DISCOVERY_DOY DISCOVERY_TIME  STAT_CAUSE_CODE STAT_CAUSE_DESCR  CONT_DATE  CONT_DOY CONT_TIME  FIRE_SIZE FIRE_SIZE_CLASS   LATITUDE   LONGITUDE  OWNER_CODE       OWNER_DESCR STATE COUNTY FIPS_CODE  FIPS_NAME                                              Shape\n",
       "0         1       1  FS-1418826                FED   FS-FIRESTAT                    FS                USCAPNF    Plumas National Forest                  0511     Plumas National Forest                    1            PNF-47      BJ8K  FOUNTAIN                    None         None    None           None         None       2005       2453403.5             33           1300              9.0    Miscellaneous  2453403.5      33.0      1730       0.10               A  40.036944 -121.005833         5.0              USFS    CA     63       063     Plumas  b'\\x00\\x01\\xad\\x10\\x00\\x00\\xe8d\\xc2\\x92_@^\\xc0...\n",
       "1         2       2  FS-1418827                FED   FS-FIRESTAT                    FS                USCAENF  Eldorado National Forest                  0503   Eldorado National Forest                   13                13      AAC0    PIGEON                    None         None    None           None         None       2004       2453137.5            133           0845              1.0        Lightning  2453137.5     133.0      1530       0.25               A  38.933056 -120.404444         5.0              USFS    CA     61       061     Placer  b'\\x00\\x01\\xad\\x10\\x00\\x00T\\xb6\\xeej\\xe2\\x19^\\...\n",
       "2         3       3  FS-1418835                FED   FS-FIRESTAT                    FS                USCAENF  Eldorado National Forest                  0503   Eldorado National Forest                   27               021      A32W     SLACK                    None         None    None           None         None       2004       2453156.5            152           1921              5.0   Debris Burning  2453156.5     152.0      2024       0.10               A  38.984167 -120.735556        13.0  STATE OR PRIVATE    CA     17       017  El Dorado  b'\\x00\\x01\\xad\\x10\\x00\\x00\\xd0\\xa5\\xa0W\\x13/^\\...\n",
       "3         4       4  FS-1418845                FED   FS-FIRESTAT                    FS                USCAENF  Eldorado National Forest                  0503   Eldorado National Forest                   43                 6      None      DEER                    None         None    None           None         None       2004       2453184.5            180           1600              1.0        Lightning  2453189.5     185.0      1400       0.10               A  38.559167 -119.913333         5.0              USFS    CA      3       003     Alpine  b'\\x00\\x01\\xad\\x10\\x00\\x00\\x94\\xac\\xa3\\rt\\xfa]...\n",
       "4         5       5  FS-1418847                FED   FS-FIRESTAT                    FS                USCAENF  Eldorado National Forest                  0503   Eldorado National Forest                   44                 7      None  STEVENOT                    None         None    None           None         None       2004       2453184.5            180           1600              1.0        Lightning  2453189.5     185.0      1200       0.10               A  38.559167 -119.933056         5.0              USFS    CA      3       003     Alpine  b'\\x00\\x01\\xad\\x10\\x00\\x00@\\xe3\\xaa.\\xb7\\xfb]\\..."
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3291b4ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = open(\"columns.txt\", \"r\")\n",
    "\n",
    "col_txt = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f56d9411",
   "metadata": {},
   "outputs": [],
   "source": [
    "openai.api_key = 'xxxxxxxxxxxxxxxxxxxxxxxxxx'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5c1d976a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "a08f17fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataAnalysisAgent:\n",
    "    def __init__(self, db_path):\n",
    "        self.connection = sqlite3.connect(db_path)\n",
    "        self.cursor = self.connection.cursor()\n",
    "        self.valid_columns = self.fetch_column_names()\n",
    "        self.valid_tables = {'fires'}\n",
    "\n",
    "    def fetch_column_names(self):\n",
    "        \"\"\" Fetches all column names from the 'fires' table to ensure query accuracy. \"\"\"\n",
    "        self.cursor.execute(\"PRAGMA table_info(fires)\")\n",
    "        return {item[1] for item in self.cursor.fetchall()}  # Set of column names\n",
    "\n",
    "\n",
    "    def is_valid_sql_columns_table(self, sql_query):\n",
    "        \"\"\" Check if the SQL is syntactically correct and if it contains valid columns and tables. \"\"\"\n",
    "        try:\n",
    "            parsed = sqlparse.parse(sql_query)[0]\n",
    "            tokens = [token for token in parsed.tokens if not token.is_whitespace]\n",
    "            print(tokens)\n",
    "\n",
    "            found_tables = set()\n",
    "            found_columns = set()\n",
    "            for token in tokens:\n",
    "                if(1):\n",
    "                    value = token.value\n",
    "                    if value in self.valid_columns:\n",
    "                        found_columns.add(value)\n",
    "                    elif value in self.valid_tables:\n",
    "                        found_tables.add(value)\n",
    "            print(found_tables)\n",
    "            print(found_columns)\n",
    "\n",
    "            if not found_tables.issubset(self.valid_tables) or found_tables == set():\n",
    "                print(\"Query contains invalid table names.\")\n",
    "                return False\n",
    "            if not found_columns.issubset(self.valid_columns) :\n",
    "                print(\"Query contains invalid column names.\")\n",
    "                return False\n",
    "            return True\n",
    "\n",
    "        except sqlparse.exceptions.SQLParseError:\n",
    "            return False\n",
    "\n",
    "    def is_valid_sql_columns(self, sql_query):\n",
    "        \"\"\" Check if the SQL is syntactically correct and if it contains valid columns. \"\"\"\n",
    "        try:\n",
    "            parsed = sqlparse.parse(sql_query)[0]\n",
    "            tokens = {token.value for token in parsed.tokens if token.ttype is sqlparse.tokens.Name}\n",
    "            if not tokens.issubset(self.valid_columns):\n",
    "                print(\"Query contains invalid column names.\")\n",
    "                return False\n",
    "            return True\n",
    "        except sqlparse.exceptions.SQLParseError:\n",
    "            return False\n",
    "\n",
    "    def is_valid_sql(self,sql_query):\n",
    "        try:\n",
    "            sqlparse.parse(sql_query)\n",
    "            return True\n",
    "        except sqlparse.exceptions.SQLParseError:\n",
    "            return False\n",
    "        \n",
    "\n",
    "\n",
    "\n",
    "    def generate_sql_query(self, user_query):\n",
    "        \"\"\"\n",
    "        Uses OpenAI's GPT-3.5 Turbo to generate an SQL query based on the user's natural language query.\n",
    "        \n",
    "        Args:\n",
    "            user_query (str): A natural language query from the user.\n",
    "\n",
    "        Returns:\n",
    "            str: SQL query string generated by the AI.\n",
    "            dict: Any SQL parameters for safe query execution (not used here for simplicity).\n",
    "        \"\"\"\n",
    "        prompt = ''' You are an assistant trained to convert natural language to SQL.\n",
    "        <INST> \n",
    "            \n",
    "        DATABASE DETAILS:\n",
    "        - Database: SQLite\n",
    "        - Table Name: fires\n",
    "        - cnx = sqlite3.connect('FPA_FOD_20170508.sqlite')\n",
    "        - USE cnx to query the database using sqlite commands\n",
    "\n",
    "\n",
    "        MEMORY:\n",
    "        Context: The 'fires' table contains a comprehensive spatial database of wildfires that occurred in the United States from 1992 to 2015. This data supports national Fire Program Analysis (FPA) systems. The records, sourced from federal, state, and local fire organizations, include essential elements such as discovery date, final fire size, and point locations precise to at least the Public Land Survey System (PLSS) section (1-square mile grid). The data conforms to standards set by the National Wildfire Coordinating Group (NWCG) and includes error-checking and redundancy removal. This data publication, known as the Fire Program Analysis fire-occurrence database (FPA FOD), comprises 1.88 million geo-referenced wildfire records, representing 140 million acres burned over a 24-year period.\n",
    "\n",
    "        Column Information:\n",
    "\n",
    "        '''\n",
    "        prompt+= col_txt\n",
    "        prompt+='''\n",
    "        <\\INST>\n",
    "        Assistant: ```sql\n",
    "        <SQL>\n",
    "        ```\n",
    "        '''\n",
    "        response = openai.ChatCompletion.create(\n",
    "            model=\"gpt-3.5-turbo\",  # Specify the appropriate model\n",
    "            messages=[\n",
    "                {\"role\": \"system\", \"content\": prompt},\n",
    "                {\"role\": \"user\", \"content\": user_query}\n",
    "            ]\n",
    "        )\n",
    "        sql_query = response['choices'][0]['message']['content']\n",
    "        pattern = r\"```sql(.*?)```\"\n",
    "        print(sql_query)\n",
    "        try:\n",
    "            # Search for the pattern in the text\n",
    "            match = re.search(pattern, sql_query, re.DOTALL)\n",
    "            sql_extracted = match.group(1).strip()\n",
    "\n",
    "            if match:\n",
    "                # Extract the SQL query, stripping any leading or trailing whitespace\n",
    "                if(self.is_valid_sql(sql_extracted)):\n",
    "                    print(\"Valid SQL query\")\n",
    "                else:\n",
    "                    print(\"Regenerate the query\") #return to base\n",
    "                if(self.is_valid_sql_columns(sql_extracted)):\n",
    "                    print(\"Valid SQL query and columns\")\n",
    "                else:\n",
    "                    print(\"Regenerate the query\") #return to base\n",
    "                return match.group(1).strip(),{}\n",
    "            else:\n",
    "                return \"\",{}\n",
    "\n",
    "            return sql_query,{}\n",
    "        except Exception as e:\n",
    "            print(f\"Failed to generate SQL query: {str(e)}\")\n",
    "            return \"\", {}\n",
    "\n",
    "#     def execute_query(self, query):\n",
    "#         return pd.read_sql_query(query, self.connection)\n",
    "    def execute_query(self, query, timeout=10, max_rows = 1000000):  # Timeout in seconds for query execution\n",
    "#         query = f\"{query} LIMIT {max_rows}\" \n",
    "        start_time = time.time()\n",
    "        try:\n",
    "            df = pd.read_sql_query(query, self.connection)\n",
    "            if time.time() - start_time > timeout:\n",
    "                raise TimeoutError(\"Query execution exceeded the time limit.\")\n",
    "            return df\n",
    "        except TimeoutError as te:\n",
    "            print(te)\n",
    "            return pd.DataFrame()  # Return empty DataFrame on timeout\n",
    "        except Exception as e:\n",
    "            print(f\"Error executing query: {str(e)}\")\n",
    "            return pd.DataFrame()\n",
    "\n",
    "    def analyze_data(self, dataframe):\n",
    "        # Perform analysis like aggregations, statistical tests, etc.\n",
    "        return dataframe.describe()\n",
    "    \n",
    "    def get_query_plan(self,query):\n",
    "        plan_query = f\"EXPLAIN {query}\"\n",
    "        return pd.read_sql_query(plan_query, self.connection)\n",
    "\n",
    "\n",
    "    def handle_query(self, user_query):\n",
    "        sql_query, params = self.generate_sql_query(user_query)\n",
    "        print(\"query\")\n",
    "        print(sql_query)\n",
    "        try:\n",
    "            data = self.execute_query(sql_query)\n",
    "            print(data)\n",
    "            if data.empty:\n",
    "                return \"No data returned.\"\n",
    "            else:\n",
    "                result = self.analyze_data(data)\n",
    "                return result\n",
    "        except Exception as e:\n",
    "            return f\"Error executing query: {str(e)}\"\n",
    "\n",
    "\n",
    "# Usage\n",
    "agent = DataAnalysisAgent('FPA_FOD_20170508.sqlite')\n",
    "\n",
    "# result = agent.handle_query(\"What is the average fire size?\")\n",
    "# print(result)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "cf9f884b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "..."
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[<DML 'SELECT' at 0x287715960>, <Wildcard '*' at 0x2862D5AE0>, <Keyword 'FROM' at 0x2862D7820>, <Identifier 'fires' at 0x2863076D0>, <Punctuation ';' at 0x2862D6BC0>]\n",
      "{'fires'}\n",
      "set()\n",
      "[<DML 'SELECT' at 0x2862D64A0>, <Wildcard '*' at 0x2862D5D80>, <Keyword 'FROM' at 0x2862D7700>, <Identifier 'unknow...' at 0x2863075D0>, <Punctuation ';' at 0x2862D5FC0>]\n",
      "set()\n",
      "set()\n",
      "Query contains invalid table names.\n",
      "```sql\n",
      "SELECT AVG(FIRE_SIZE) AS Average_Fire_Size\n",
      "FROM fires;\n",
      "```\n",
      "Valid SQL query\n",
      "Valid SQL query and columns\n",
      "query\n",
      "SELECT AVG(FIRE_SIZE) AS Average_Fire_Size\n",
      "FROM fires;\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "."
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   Average_Fire_Size\n",
      "0          74.520158\n",
      "result\n",
      "       Average_Fire_Size\n",
      "count           1.000000\n",
      "mean           74.520158\n",
      "std                  NaN\n",
      "min            74.520158\n",
      "25%            74.520158\n",
      "50%            74.520158\n",
      "75%            74.520158\n",
      "max            74.520158\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      ".\n",
      "----------------------------------------------------------------------\n",
      "Ran 5 tests in 1.672s\n",
      "\n",
      "OK\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "```sql\n",
      "SELECT AVG(FIRE_SIZE)\n",
      "FROM fires;\n",
      "```\n",
      "Valid SQL query\n",
      "Valid SQL query and columns\n"
     ]
    }
   ],
   "source": [
    "#TESTING\n",
    "\n",
    "import unittest\n",
    "\n",
    "class TestDataAnalysisAgent(unittest.TestCase):\n",
    "\n",
    "    def setUp(self):\n",
    "        self.agent = DataAnalysisAgent('FPA_FOD_20170508.sqlite')\n",
    "\n",
    "    def test_fetch_column_names(self):\n",
    "        columns = self.agent.fetch_column_names()\n",
    "        self.assertIsInstance(columns, set)\n",
    "        self.assertTrue('FIRE_SIZE' in columns)\n",
    "\n",
    "    def test_is_valid_sql(self):\n",
    "        valid_sql = \"SELECT * FROM fires;\"\n",
    "        invalid_sql = \"SELECT * FROM unknown_table;\"\n",
    "        self.assertTrue(self.agent.is_valid_sql_columns_table(valid_sql))\n",
    "        self.assertFalse(self.agent.is_valid_sql_columns_table(invalid_sql))\n",
    "\n",
    "    def test_sql_query_generation(self):\n",
    "        response = self.agent.generate_sql_query(\"average fire size\")\n",
    "        self.assertNotEqual(response, (\"\", {}))  # Assuming non-empty returns on valid prompts\n",
    "\n",
    "    def test_execute_query(self):\n",
    "        query = \"SELECT * FROM fires LIMIT 10;\"\n",
    "        result = self.agent.execute_query(query)\n",
    "        self.assertFalse(result.empty)\n",
    "\n",
    "    def test_query_handling(self):\n",
    "        # Test the end-to-end process of handling a query\n",
    "        result = self.agent.handle_query(\"What is the average fire size?\")\n",
    "        print(\"result\")\n",
    "        print(str(result))\n",
    "        self.assertNotEqual(str(result), \"\")\n",
    "        self.assertIn('mean', str(result))  # Assuming 'mean' would be part of the description output\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    unittest.main(argv=['first-arg-is-ignored'], exit=False)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "1d35bd28",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f5923f8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
